package com.yanggu.flink.datastream_api.window

import cn.hutool.core.date.DateUtil
import org.apache.flink.api.common.ExecutionConfig
import org.apache.flink.api.common.typeutils.TypeSerializer
import org.apache.flink.streaming.api.environment
import org.apache.flink.streaming.api.scala.{StreamExecutionEnvironment, _}
import org.apache.flink.streaming.api.windowing.assigners.WindowAssigner
import org.apache.flink.streaming.api.windowing.time.Time
import org.apache.flink.streaming.api.windowing.triggers.{EventTimeTrigger, Trigger}
import org.apache.flink.streaming.api.windowing.windows.TimeWindow

import java.util
import java.util.{Collections, Date}

/**
 * 实现根据周或月的窗口划分窗口
 * 或者每个月第一天的00:00:00到最后一天的23:59:59
 * 实现参考了 TumblingEventTimeWindows
 *
 * @param tag 标签 month  or  week
 * @tparam T 需要划分窗口的数据类型(输入类型)
 */
class CustomTumblingEventTimeWindowAssigner[T](tag: String) extends WindowAssigner[T, TimeWindow] {
  //窗口分配的主要方法，需要为每一个元素指定所属的分区
  override def assignWindows(element: T, timestamp: Long, context: WindowAssigner.WindowAssignerContext): util.Collection[TimeWindow] = {
    val offset: (Long, Long) = tag match {
      case "week" => getTimestampFromWeek(timestamp)
      case "month" => getTimestampFromMonth(timestamp)
      case "quarter" => getTimestampFromQuarter(timestamp)
      case "year" => getTimestampFromYear(timestamp)
    }
    //分配窗口
    Collections.singletonList(new TimeWindow(offset._1, offset._2))
  }

  //注意此处需要进行类型的转换，否则或编译出错，java版本好像没问题，但是java对于上面的offset处理有点难搞，所以放弃了
  override def getDefaultTrigger(env: environment.StreamExecutionEnvironment): Trigger[T, TimeWindow] = EventTimeTrigger.create().asInstanceOf[Trigger[T, TimeWindow]]

  override def getWindowSerializer(executionConfig: ExecutionConfig): TypeSerializer[TimeWindow] = new TimeWindow.Serializer()

  //是否使用事件时间
  override def isEventTime: Boolean = true

  /**
   * 周窗口
   */
  def getTimestampFromWeek(timestamp: Long): (Long, Long) = {
    (DateUtil.beginOfWeek(new Date(timestamp)).getTime, DateUtil.endOfWeek(new Date(timestamp)).getTime + 1)
  }

  /**
   * 月窗口
   */
  def getTimestampFromMonth(timestamp: Long): (Long, Long) = {
    (DateUtil.beginOfMonth(new Date(timestamp)).getTime, DateUtil.endOfMonth(new Date(timestamp)).getTime + 1)
  }

  /**
   * 季度窗口
   */
  def getTimestampFromQuarter(timestamp: Long): (Long, Long) = {
    (DateUtil.beginOfQuarter(new Date(timestamp)).getTime, DateUtil.endOfQuarter(new Date(timestamp)).getTime + 1)
  }

  /**
   * 年窗口
   */
  def getTimestampFromYear(timestamp: Long): (Long, Long) = {
    (DateUtil.beginOfYear(new Date(timestamp)).getTime, DateUtil.endOfYear(new Date(timestamp)).getTime + 1)
  }

}

object Test {

  def main(args: Array[String]): Unit = {
    println(DateUtil.formatDateTime(DateUtil.endOfYear(new Date())))
    test1(1)
    test1(7)
    val tuple = new CustomTumblingEventTimeWindowAssigner[Object]("week").getTimestampFromWeek(System.currentTimeMillis())
    println(s"week start: ${DateUtil.formatDateTime(new Date(tuple._1))}, end: ${DateUtil.formatDateTime(new Date(tuple._2))}")
    val tuple1 = new CustomTumblingEventTimeWindowAssigner[Object]("month").getTimestampFromMonth(System.currentTimeMillis())
    println(s"month start: ${DateUtil.formatDateTime(new Date(tuple1._1))}, end: ${DateUtil.formatDateTime(new Date(tuple1._2))}")
    val tuple2 = new CustomTumblingEventTimeWindowAssigner[Object]("quarter").getTimestampFromQuarter(System.currentTimeMillis())
    println(s"quarter start: ${DateUtil.formatDateTime(new Date(tuple2._1))}, end: ${DateUtil.formatDateTime(new Date(tuple2._2))}")
    val tuple3 = new CustomTumblingEventTimeWindowAssigner[Object]("year").getTimestampFromYear(System.currentTimeMillis())
    println(s"year start: ${DateUtil.formatDateTime(new Date(tuple3._1))}, end: ${DateUtil.formatDateTime(new Date(tuple3._2))}")
  }

  private def test1(day: Long) = {
    val oneDayMills = Time.days(day).toMilliseconds
    val start = TimeWindow.getWindowStartWithOffset(System.currentTimeMillis(), -Time.hours(8L).toMilliseconds, oneDayMills)
    val end = start + oneDayMills
    println(s"start: ${DateUtil.formatDateTime(new Date(start))}, end: ${DateUtil.formatDateTime(new Date(end))}")
  }
}